﻿#pragma once
#include <onnxruntime_cxx_api.h>
#include <opencv2/opencv.hpp>
#include <vector>
#include <string>
#include <memory>
//#include <cuda_provider_factory.h>

class OnnxDetector
{
public:
    std::string modelPath;//图片路径
    const int INPUT_SIZE = 640; // YOLOv8 标准输入尺寸
    float CONFIDENCE_THRESHOLD = 0.25f; // 置信度阈值
    float IOU_THRESHOLD = 0.45f; // IoU阈值(NMS使用)
    const std::vector<std::string> CLASS_NAMES;// 类别名称

    // 检测结果结构体
    struct Detection {
        cv::Rect_<float> bbox;
        float confidence;
        int class_id;
        std::string label;
    };

    // 转换参数结构体
    struct TransformParams {
        float scale;
        int offset_x;
        int offset_y;
        int scaled_width;
        int scaled_height;
    };

    OnnxDetector(const std::string& modelPath, const std::vector<std::string>& class_names, float confidenceThreshold = 0.5f, float nmsThreshold = 0.45f);

    ~OnnxDetector();

    std::vector<Detection> detect(const cv::Mat& image);
    cv::Mat draw_detections(const cv::Mat& image, const std::vector<Detection>& detections);
    void save_detections_to_file(const std::vector<Detection>& detections, const std::string& filename);
    void init();

private:
    // ONNX Runtime环境
    std::unique_ptr<Ort::Env> env_;
    std::unique_ptr<Ort::Session> session_;
    Ort::SessionOptions session_options_;

    std::tuple<cv::Mat, TransformParams> preprocess_image(const cv::Mat& image);
    std::vector<Detection> parse_model_output(const cv::Mat& input_tensor, const TransformParams& params, const cv::Mat& original_image);
    std::vector<Detection> apply_nms(std::vector<Detection>& detections);
    float calculate_iou(const cv::Rect_<float>& rect1, const cv::Rect_<float>& rect2);
    cv::Scalar get_class_color(int class_id);
};
